import copy
import time
import os
import math
import yaml
import urx
import math3d as m3d
import numpy as np
from numpy import cos, sin
from algorithm.mathlib import utilities as mathutl
from algorithm.mathlib.utilities import PoseRobot
from algorithm.dataprocessing import utilities as dputl
from pyquaternion import Quaternion
from vrep_api import vrep
import FTsensor as FT
# import utils

# execude mode
sim_mode = 0
real_mode = 1
# filter_type
no_filter = 0
lowpass_filter = 1
class RobotUR5(object):
    def __init__(self):

        # self.robot = urx.Robot("192.168.1.100")
        self.init_pos = np.array([-0.52694, -0.04607, 0.63414])
        self.init_ori = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
        # equal to init_ori
        self.init_qtn = Quaternion(np.array([0, 1, 0, 0]))
        self.pose_des = PoseRobot()
        self.pose_curr = PoseRobot()

        main_path = os.path.dirname((os.path.dirname(os.path.realpath(__file__))))
        yaml_file = os.path.join(main_path, "config/RobotConfig.yaml")
        with open(yaml_file, encoding='utf-8') as f:
            content = f.read()
        self.yaml_data = yaml.load(content, Loader=yaml.FullLoader)
        self.pose_des.PoseInit(np.array(self.yaml_data["pose_desire"]))

        self.ft_sensor = FT.FTSensor(lowpass_filter)
        self.controller= None

    def set_ft_sensor(self):
        self.ft_sensor._Foffset = self.yaml_data["force_offset"]
        self.ft_sensor._Toffset = self.yaml_data["torque_offset"]
        self.ft_sensor._Gravity = self.yaml_data["gravity"]
        self.ft_sensor._Pcom = self.yaml_data["tool_pos_com"]
        self.ft_sensor.wrench_safety = self.yaml_data["wrench_safety"]
        self.ft_sensor.PoseTool2Sensor = self.yaml_data["transform_tool_to_sensor"]

        self.ft_sensor._Pcom = self.sim_get_pos_toolcom2sensor()
        self.ft_sensor._PoseTool2Sensor = self.sim_get_pose_tip2sensor()

        self._adj_tool_2_sensor = mathutl.SE32Adj(mathutl.posem2SE3(self.ft_sensor.PoseTool2Sensor))


    def set_init_pose(self, acc=0.01, vel=0.01):
        trans = m3d.Transform()
        trans.orient = self.init_ori
        trans.pos = self.init_pos
        if self.ctrlmode == real_mode:
            self.robot.set_pose(trans, acc=acc, vel=vel)

    def sim_send_data(self):
        """设定机械臂末端到达位置position"""
        return self.controller.core

    def sim_get_pose_tip2sensor(self):
        sim_ret, pos_tip2sensor = vrep.simxGetObjectPosition(self.sim_client, self.sim_tip_handle, self.sim_sensor_frame_handle, vrep.simx_opmode_blocking)
        sim_ret, ori_tip2sensor = vrep.simxGetObjectOrientation(self.sim_client, self.sim_tip_handle, self.sim_sensor_frame_handle, vrep.simx_opmode_blocking)
        pose_tip2sensor = np.concatenate(pos_tip2sensor, ori_tip2sensor)
        return pose_tip2sensor

    def sim_get_pos_toolcom2sensor(self):
        sim_ret, pos_toolcom2sensor = vrep.simxGetObjectPosition(self.sim_client, self.sim_tool_frame_handle, self.sim_sensor_frame_handle, vrep.simx_opmode_blocking)
        return pos_toolcom2sensor

    def sim_get_ftdata(self):
        sim_ret, state, forceVector, torqueVector = vrep.simxReadForceSensor(self.sim_client,self.sim_ft_handle,vrep.simx_opmode_streaming)
        ftdata = np.concatenate(forceVector, torqueVector)
        return ftdata

    def get_ftdata(self):
        return self.ft_sensor.getWrenchNetTool(self.pose_curr.qtn)

    def sim_read_data(self, data_to_read):
        self.ft_sensor.new_force = data_to_read.forceVector
        self.ft_sensor.new_torque = data_to_read.torqueVector
        self.pose_curr.pos = data_to_read.tip_position
        self.pose_curr.qtn = mathutl.xyz2qtn(data_to_read.tip_orientation)


    def my_set_pose(self, theta, position, acc, vel):
        """设定机械臂末端到达位置position"""
        trans = m3d.Transform()
        orientation = np.array([[cos(theta), sin(theta), 0], [sin(theta), -cos(theta), 0], [0, 0, -1]])
        trans.orient = orientation
        trans.pos = position
        if self.ctrlmode == real_mode:
            self.robot.set_pose(trans, acc=acc, vel=vel)